package com.bw.test1;

import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.regression.LinearRegTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.stream.StreamOperator;
import com.alibaba.alink.operator.stream.regression.LinearRegPredictStreamOp;
import com.alibaba.alink.operator.stream.source.MemSourceStreamOp;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class LinearRegPredictStreamOpTest {
    @Test
    public void testLinearRegPredictStreamOp() throws Exception {
        List<Row> df = Arrays.asList(
                Row.of(2, 1, 1),
                Row.of(3, 2, 1),
                Row.of(4, 3, 2),
                Row.of(2, 4, 1),
                Row.of(2, 2, 1),
                Row.of(4, 3, 2),
                Row.of(1, 2, 1)
        );
        BatchOperator<?> batchData = new MemSourceBatchOp(df, "f0 int, f1 int, label int");
        StreamOperator<?> streamData = new MemSourceStreamOp(df, "f0 int, f1 int, label int");
        String[] colnames = new String[]{"f0", "f1"};
        // 用的批的训练
        BatchOperator<?> lr = new LinearRegTrainBatchOp()
                .setFeatureCols(colnames)
                .setLabelCol("label");
        BatchOperator<?> model = batchData.link(lr);

        //用的预测
        StreamOperator<?> predictor = new LinearRegPredictStreamOp(model)
                .setPredictionCol("pred");
        predictor.linkFrom(streamData).print();

        // 流必须要执行
        StreamOperator.execute();
    }
}